from flask import Blueprint, request, jsonify, current_app, stream_with_context,Response
import logging, os
from utils import common
chat_bp = Blueprint('chat', __name__, url_prefix='/chat')

from ._impl import get_response

import tiktoken
enc = tiktoken.get_encoding("cl100k_base")


# modules_and_functions = common.find_modules_and_functions(os.path.dirname(__file__))

# for module, impls in modules_and_functions:
#     for impl in impls:
#         admin_bp.add_url_rule(f"{module.__name__}/{impl if impl.find('_') == -1  else '/'.join(impl.split('_'))}", impl, getattr(module, impl))

# chat_bp = common.add_url_to_bp(os.path.dirname(__file__), chat_bp)

@chat_bp.route('/list')
def list():
    params = request.args
    result = current_app.config['DB'].list_chat(params.get('user_id'))

    return jsonify({"code":0, "data":result})

@chat_bp.route('/history')
def history():
    params = request.args
    result = current_app.config['DB'].get_chat_history(params.get('chat_id'))
    return jsonify({"code":0, "data":result})

@chat_bp.route('/tokens', methods=['POST'])
def tokens():
    params = request.json
    if params.get('text'):
        return jsonify({'code':0, 'data': len(enc.encode(params.get('text').replace('</p>', '').replace('\n', '').replace('<p>', '')))})
    
    ## 获取最近模型回复
    result = current_app.config['DB'].get_tokens(params)
    print(result)
    return jsonify({"code":0, "data":result[0][0]})

@chat_bp.route('/new', methods=['GET'])
def new():
    # 实现sse协议。流式对话输出
    params = dict(request.args)#request.json
    ## 检查用户状态
    response = Response(stream_with_context(get_response(params)), mimetype="text/event-stream")
    response.headers['Content-Type'] = "text/event-stream"
    return response
    
@chat_bp.route('/update', methods=['POST'])
def update():
    params = request.json
    ## 检查用户状态
    current_app.config['DB'].update_chat_title(params)
    return jsonify({'code':0})

@chat_bp.route('/remove', methods=['GET'])
def remove():
    params = dict(request.args)
    ## 检查用户状态
    current_app.config['DB'].remove_chat_by_id(params)
    return jsonify({'code':0})
        
@chat_bp.route('/clear', methods=['GET'])
def clear():
    params = dict(request.args)
    ## 检查用户状态
    current_app.config['DB'].clear_chat_by_id(params)
    return jsonify({'code':0})

@chat_bp.record
def log(app):
    logging.info('register chat blueprint')
